#include "numuCC1P1PiSelection.hxx"
#include "numuCCSelection.hxx"
#include "baseSelection.hxx"
#include "CutUtils.hxx"

numuCC1P1PiSelection::numuCC1P1PiSelection(bool forceBreak): SelectionBase(forceBreak) 
{
  // define the steps
  DefineSteps();  
}

void numuCC1P1PiSelection::DefineSteps()
{
    // Cuts must be added in the right order
    // last "true" means the step sequence is broken if cut is not passed (default is "false")
    // Add inclusive CC steps to all branches

    AddStep(StepBase::kCut,    "event quality", new EventQualityCut(), true);
    AddStep(StepBase::kCut,    "> 0 tracks ", new TotalMultiplicityCut(), true); // Greater than Zero tracks with a least one in the TPC (I assume TPC2 here) - !!find out!!
    AddStep(StepBase::kAction, "find leading tracks", new FindLeadingTracksAction());  // From selected tracks - find quality tracks in FGD1 FV only
    AddStep(StepBase::kAction, "find vertex", new FindVertexAction());  // Determine vertex of interaction in given detector (FGD1 here), which is the start position of the HM Track
    AddStep(StepBase::kCut,    "quality+fiducial", new TrackQualityFiducialCut(), true);// Selects good quality tracks originated from FGD1
    AddStep(StepBase::kAction, "find veto track", new FindVetoTrackAction()); // Find next highest momentum track after muon candidate (named vetotrack) - Is this needed YES accounts for tracks coming from behind the muons start point
    AddStep(StepBase::kCut,    "veto", new ExternalVetoCut()); // Removes events entering FGD1 FV/behind muon
    AddStep(StepBase::kAction, "find oofv track", new FindOOFVTrackAction());// Get track out of FGD1 fudicial volume
    AddStep(StepBase::kCut,    "muon PID", new MuonPIDCut());// Determine muon PID given muon canditate - Uses HMN track
    AddStep(StepBase::kAction, "find michel e-", new FindMichelElectronsAction()); // Find any michel electrons
    AddStep(StepBase::kCut,    "no michel ele", new MichelElectronCut()); // Cut on any michel electrons found
    AddStep(StepBase::kAction, "find tracker tracks", new FindTrackerTracksAction()); //local action; on std_numucc1P1Pi_actions::FindTracksWithTPC,!!Not sure what this does !!
    //From here are my own cuts: 1st cut Track Charge, 2nd cut FindCandidates (Pion, Proton), 3 point common vertex
    AddStep(StepBase::kCut, "Select 3 Charged Tracks (2pos, 1neg)", new TrackChargeCut());
    AddStep(StepBase::kCut, "Proton & Pion PID", new ParticlePIDCut());
    AddStep(StepBase::kCut, "common vertex", new CommonVertexCut()); // Select the common vertex for the tracks

    //AddStep(StepBase::kAction, "Find correct particle sense", new FindCorrectParticleSenseAction());
    
    //very important! without it no events shown in Scan
    SetBranchAlias(0,"muTPC-pTPC"); //What is this - TBranch.
}

bool numuCC1P1PiSelection::FillEventSummary(const AnaBoxB& box, AnaEventB& event, Int_t allCutsPassed[])
{
    (void)box;

    // Should set event sample enumeration here
    if(allCutsPassed[0]) event.Summary->EventSample = gkSampleEnumCC1P1Pi;
    return (event.Summary->EventSample != nd280Samples::kUnassigned);
}

bool FindMichelElectronsAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    cutUtils::FindFGDMichelElectrons(event,box);
    
    return true;
}

bool MichelElectronCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    // Temporarily, a Float_t cut to make sure it works with OldFlatTree, FlatTree and oaAnalysis
    if (box.nFGD1MichelElectrons == 0)
        return true;
    else
        return false;
}


bool FindTrackerTracksAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    AnaBoxCC1P1Pi * cc1P1Pibox = static_cast<AnaBoxCC1P1Pi*>(&box);
    
    // all Tracker tracks
    cc1P1Pibox->TrackerTracks.clear();
    
    std_numucc1P1Pi_actions::FindTracksWithTPC(event,*cc1P1Pibox, true, SubDetId::kFGD1);
    //needed!! otherwise the selected sample is different!! MAYBE due to michel electron cut
    std_numucc1P1Pi_actions::FindTracksWithFGDAndNoTPC(event, *cc1P1Pibox, SubDetId::kFGD1);
    
    return true;
}

bool TrackChargeCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;
    
    AnaBoxCC1P1Pi * cc1P1Pibox = static_cast<AnaBoxCC1P1Pi*>(&box);
    
    // one negative track with good quality and in FV
    if (box.nNegativeTPCtracks==1 && box.nPositiveTPCtracks==2 && cc1P1Pibox->TrackerTracks.size() == 3)
        return true;
    else
        return false;
}

bool ParticlePIDCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;
    
    AnaBoxCC1P1Pi * cc1P1Pibox = static_cast<AnaBoxCC1P1Pi*>(&box);
    
    for(int i=0; i<box.nPositiveTPCtracks; i++){
        AnaTrackB *ptrack = box.PositiveTPCtracks[i];
        
        if(!cutUtils::FiducialCut(*ptrack,SubDetId::kFGD1)) continue;//Make sure the track originates from FGD1
        
        double ElectrLklh = anaUtils::GetPIDLikelihood(*ptrack,1);//Electron
        double ProtonLklh = anaUtils::GetPIDLikelihood(*ptrack,2);//Proton
        double PionLklh = anaUtils::GetPIDLikelihood(*ptrack,3);//Pion
        double norm = ElectrLklh + ProtonLklh + PionLklh;
        ProtonLklh /= norm;
        ElectrLklh /= norm;
        PionLklh /= norm;
        //printf("ProtonLk: %f, ElectronLk: %f, PionLk: %f\n",ProtonLklh, ElectrLklh, PionLklh);
        
       /* if(ProtonLklh > PionLklh && ProtonLklh > ElectrLklh){
            cc1P1Pibox->ProtonCandidate = ptrack;
          //  printf("Proton track selected\n");
        }
        else if(PionLklh > ProtonLklh && PionLklh > ElectrLklh){
            cc1P1Pibox->PionCandidate = ptrack;
            //printf("Pion track selected\n");
        }*/
        if(std_numucc1P1Pi_cuts::IsProton(ptrack) && !std_numucc1P1Pi_cuts::IsPion(ptrack)){
            cc1P1Pibox->ProtonCandidate = ptrack;
            //printf("Proton track selected\n");
        }
        else if(std_numucc1P1Pi_cuts::IsPion(ptrack) && !std_numucc1P1Pi_cuts::IsProton(ptrack)){
            cc1P1Pibox->PionCandidate = ptrack;
            //printf("Pion track selected\n");
        }
    }

    if(cc1P1Pibox->PionCandidate && cc1P1Pibox->ProtonCandidate){
        cc1P1Pibox->MuonCandidate = box.HMNtrack;
        //printf("Proton and pion candidates selected\n");
        return true;
    }
    else return false;
}
/*
bool FindCorrectParticleSenseAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    //Determine both pion and proton sense is correct - particle/antiparticle comes from vertex: Corrects for start/finish position of particle in detector (Recon failures)
    (void)event;
    
    AnaBoxCC1P1Pi * cc1P1Pibox = static_cast<AnaBoxCC1P1Pi*>(&box);
    //Proton Candidate
    if (cc1P1Pibox->ProtonCandidate)
        // it must reverse the track only if it is a FGD track (no TPC segments)
        if (!anaUtils::TrackUsesDet(*(cc1P1Pibox->ProtonCandidate), SubDetId::kTPC))
            std_numucc1P1Pi_actions::FindCorrectSense(cc1P1Pibox->ProtonCandidate, box.Vertex);
    //Pion Candidate
    if (cc1P1Pibox->PionCandidate)
        // it must reverse the track only if it is a FGD track (no TPC segments)
        if (!anaUtils::TrackUsesDet(*(cc1P1Pibox->PionCandidate), SubDetId::kTPC))
            std_numucc1P1Pi_actions::FindCorrectSense(cc1P1Pibox->PionCandidate, box.Vertex);
    
    return true;
}*/

bool CommonVertexCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;
    
    AnaBoxCC1P1Pi * cc1P1Pibox = static_cast<AnaBoxCC1P1Pi*>(&box);
   
    if(cc1P1Pibox->MuonCandidate) printf("MuonCandidate Yes\n");
    if(cc1P1Pibox->ProtonCandidate) printf("ProtonCandidate Yes\n");
    if(cc1P1Pibox->PionCandidate) printf("PionCandidate Yes\n");
    
    if (!cc1P1Pibox->MuonCandidate || !cc1P1Pibox->ProtonCandidate || !cc1P1Pibox->PionCandidate){
        //printf("failed full candidates selection");
        return false;
    }
    
    const bool kpass1 = cutUtils::CommonVertexCut(*(cc1P1Pibox->MuonCandidate), *(cc1P1Pibox->ProtonCandidate),_cutX, _cutY, _cutZ);
    const bool kpass2 = cutUtils::CommonVertexCut(*(cc1P1Pibox->PionCandidate), *(cc1P1Pibox->ProtonCandidate),_cutX, _cutY, _cutZ);
    const bool kpass3 = cutUtils::CommonVertexCut(*(cc1P1Pibox->MuonCandidate), *(cc1P1Pibox->PionCandidate),_cutX, _cutY, _cutZ);
    if(kpass1 && kpass2 && kpass3) printf("Event Selected\n");

    return (kpass1 && kpass2 && kpass3);
}

//Xianguo's PID Code
///////////////////////////////////////////////////
Bool_t std_numucc1P1Pi_cuts::IsProton(AnaTrackB * track)
{
    if(track && anaUtils::GetPIDLikelihood(*(track),2)>0.5)
        return kTRUE;
    else
        return kFALSE;
}

//copied from numuCCMultiPiUtils::TPCpionSelection
Bool_t std_numucc1P1Pi_cuts::IsPion(AnaTrackB * track)
{
    if(!track)
        return kFALSE;
    
    if ( anaUtils::GetPIDLikelihood(*track,3) < 0.3) return false;
    
    double cut1 = (anaUtils::GetPIDLikelihood(*track,0)+anaUtils::GetPIDLikelihood(*track,3))/(1.-anaUtils::GetPIDLikelihood(*track,2));
    
    if( track->Momentum < 500.  && cut1 < 0.8 ) return false;
    
    return true;
}
///////////////////////////////////////////////////

void std_numucc1P1Pi_actions::FindTracksWithTPC(const AnaEventB& event, AnaBoxCC1P1Pi& cc1P1Pibox, bool qualityCut, SubDetId::SubDetEnum det)
{
  std::vector<AnaTrackB*> allTracks[5];

  // all tracks using TPC
  AnaTrackB* selTracks[100];
  int nTPC = anaUtils::GetAllTracksUsingTPC(event, selTracks);

  //loop over tpc tracks
  for (Int_t i=0;i<nTPC; ++i){
    AnaTrackB* track = selTracks[i];

    // all TPC tracks
    allTracks[4].push_back(track);

    // all Tracker tracks
    cc1P1Pibox.TrackerTracks.push_back(track);

    //check if the track starts in FGD1 volume
    if (anaUtils::InDetVolume(SubDetId::kFGD1,track->PositionStart)) {
      if (qualityCut){
        //apply the quality cut: good quality tracks
        if (cutUtils::TrackQualityCut(*track)) { 
          // Apply the fiducial cut: save tracks out FV
          if (!cutUtils::FiducialCut(*track, det)) {

            if (track->Charge == -1)
              allTracks[0].push_back(track);
            else
              allTracks[1].push_back(track);
          }
        }
        //bad quality tracks
        else {
          // Apply the fiducial cut: save tracks in FV
          if (cutUtils::FiducialCut(*track, det))
            allTracks[2].push_back(track);
          // save tracks out FV
          else
            allTracks[3].push_back(track);
        }
      }
    }
  }

  // Fill the box
  cc1P1Pibox.NegativeTPCtracksOutFV = allTracks[0];
  cc1P1Pibox.PositiveTPCtracksOutFV = allTracks[1];
  cc1P1Pibox.BadQuaTPCtracksInFV = allTracks[2];
  cc1P1Pibox.BadQuaTPCtracksOutFV = allTracks[3];
  cc1P1Pibox.TPCtracks = allTracks[4];
}

void std_numucc1P1Pi_actions::FindTracksWithFGDAndNoTPC(const AnaEventB& event, AnaBoxCC1P1Pi& cc1P1Pibox, SubDetId::SubDetEnum det)
{
  std::vector<AnaTrackB*> allTracks[7];

  // Get all tracks with FGD and no TPC
  AnaTrackB* allFgdTracks[100];
  int nFgdTracks= anaUtils::GetAllTracksUsingFGDAndNoTPC(event, allFgdTracks);

  //loop over tracks
  for (int i=0;i<nFgdTracks;i++){
    AnaTrackB* track = allFgdTracks[i];

    // All tracks using FGD and no TPC
    allTracks[6].push_back(track);      

    // All Tracker tracks
    cc1P1Pibox.TrackerTracks.push_back(track);

    // Tracks using FGD1 (and no TPC)
    if (anaUtils::TrackUsesDet(*track,SubDetId::kFGD1)) 
      allTracks[5].push_back(track);      

    //check if the track starts in FGD1 volume
    if (anaUtils::InDetVolume(SubDetId::kFGD1,track->PositionStart)) {
      allTracks[0].push_back(track);      

      // in FV tracks
      if (cutUtils::FiducialCut(*track, det)) {
        // contained FGD tracks
        if (anaUtils::TrackUsesOnlyDet(*track, SubDetId::kFGD1))
          allTracks[1].push_back(track);
        // not-contained FGD tracks
        else
          allTracks[2].push_back(track);      
      }
      // out FV tracks
      else {
        // contained FGD tracks
        if (anaUtils::TrackUsesOnlyDet(*track, SubDetId::kFGD1))
          allTracks[3].push_back(track);
        // not-contained FGD tracks
        else 
          allTracks[4].push_back(track);      
      }
    }
  }

  cc1P1Pibox.FGDtracksInFGD1Volume = allTracks[0];
  cc1P1Pibox.FGDContTracksInFV = allTracks[1];
  cc1P1Pibox.FGDNotContTracksInFV = allTracks[2];
  cc1P1Pibox.FGDContTracksOutFV = allTracks[3];
  cc1P1Pibox.FGDNotContTracksOutFV = allTracks[4];
  cc1P1Pibox.FGDtracksCrossingFGD1 = allTracks[5];
  cc1P1Pibox.FGDtracks = allTracks[6];
}
/*
void std_numucc1P1Pi_actions::FindCorrectSense(AnaTrackB* track, AnaVertexB* vertex) 
{
  if (!track) return;
  if (!vertex) return;

  Float_t distStart = cutUtils::GetSeparationSquared(track->PositionStart, vertex->Position);
  Float_t distEnd   = cutUtils::GetSeparationSquared(track->PositionEnd, vertex->Position);

  if (distEnd < distStart){
    // reverse the track
    Float_t posTemp[4] = {track->PositionStart[0], track->PositionStart[1], track->PositionStart[2], track->PositionStart[3]};
    //    TVector3       dirTemp = track->DirectionStart;
    cutUtils::AssignArray(track->PositionEnd, track->PositionStart, 4);
    cutUtils::AssignArray(posTemp, track->PositionEnd, 4);
    //    track->DirectionStart   = - track->DirectionEnd;
    //    track->DirectionEnd     = - dirTemp;    
    track->Charge           *= -1;    
  }

}*/
